from common_import import *
from prophet import Prophet
import statsmodels.api as sm
from sklearn.linear_model import LinearRegression
import matplotlib.cm as cm
import matplotlib.colors as mcolors


def log_log_demand_model_show(data):

    # 1. 取出数据
    dates = np.array(data["date"])
    quantities = np.array(data["total_quantity"])
    sale_price_indices = np.array(data["sale_price_index"])

    # 2. 对 quantity 和 sale_price_index 取自然对数
    ln_quantities = np.log(quantities)
    ln_sale_price_indices = np.log(sale_price_indices)

    # 3. 构建线性回归模型
    X = ln_sale_price_indices
    X = sm.add_constant(X)  # 添加常数项
    y = ln_quantities

    model = sm.OLS(y, X).fit()

    # 4. 打印回归结果和参数
    print("回归方程:")
    print(
        f"ln(Quantity) = {model.params[0]:.4f} + {model.params[1]:.4f} * ln(Sale Price Index)"
    )
    print("\n回归模型的参数：")
    print(f"截距 (Intercept): {model.params[0]}")
    print(f"回归系数 (Coefficient for ln(Sale Price Index)): {model.params[1]}")
    print(f"R^2: {model.rsquared}")
    print(f"P-value: {model.pvalues[1]}")

    # 5. 可视化回归结果
    plt.figure(figsize=(10, 6))
    plt.scatter(ln_sale_price_indices, ln_quantities, label="Observed data")
    plt.plot(
        ln_sale_price_indices, model.fittedvalues, color="red", label="Fitted line"
    )
    plt.xlabel("ln(价格)")
    plt.ylabel("ln(需求)")
    plt.title("双对数模型")
    plt.legend()
    # mydraw.show_or_print("loglog.png")
    plt.show()

    return model.params[1], model.rsquared, model.pvalues[1]


def log_log_demand_get_param(data):
    # 1. 取出数据
    dates = np.array(data["date"])
    quantities = np.array(data["total_quantity"])
    sale_price_indices = np.array(data["sale_price_index"])

    # 2. 对 quantity 和 sale_price_index 取自然对数
    ln_quantities = np.log(quantities)
    ln_sale_price_indices = np.log(sale_price_indices)

    # 3. 构建线性回归模型
    X = ln_sale_price_indices
    X = sm.add_constant(X)  # 添加常数项
    y = ln_quantities
    model = sm.OLS(y, X).fit()
    return model.params[1]


def filter_data(data):
    # 提取月份（假设date格式为 YYYY-MM-DD）
    months = np.array([int(str(date).split("-")[1]) for date in data["date"]])

    # 筛选条件: 月份是6或7，且quantities不为0
    condition = (months == 6) | (months == 7)
    condition = condition & (data["quantity"] > 0)

    # 筛选数据
    filtered_data = data[condition]

    return filtered_data


def filter_data2(data):
    weekdays = np.array(
        [
            datetime.datetime.strptime(str(date), "%Y-%m-%d").weekday()
            for date in data["date"]
        ]
    )

    # 筛选条件: 星期六 (5表示星期六)，且 quantities 不为 0
    condition = (weekdays == 5) | (weekdays == 6)
    condition = condition & (data["quantity"] > 0)

    # 筛选数据
    filtered_data = data[condition]

    return filtered_data


def linear_demand_get_param(data):
    # 提取数据
    quantities = np.array(data["quantity"]).reshape(-1, 1)
    sale_price_indices = np.array(data["sales_price"]).reshape(-1, 1)
    dates = np.array(
        [datetime.datetime.strptime(str(date), "%Y-%m-%d") for date in data["date"]]
    )
    months = np.array([date.month for date in dates])

    # 构建线性回归模型
    model = LinearRegression().fit(sale_price_indices, quantities)

    # 获取唯一的月份
    unique_months = np.unique(months)

    # 生成与月份数量相同的颜色
    colors = cm.rainbow(np.linspace(0, 1, len(unique_months)))
    color_map = dict(zip(unique_months, colors))

    # 绘制散点图并根据月份标注颜色
    plt.figure(figsize=(8, 6))
    for month in unique_months:
        plt.scatter(
            sale_price_indices[months == month],
            quantities[months == month],
            color=color_map[month],
            label=f"Month {month}",
        )

    # 绘制回归线
    plt.plot(
        sale_price_indices,
        model.predict(sale_price_indices),
        color="red",
        label="Regression line",
    )

    # 添加标题和标签
    plt.title("Scatter Plot of Quantities vs Sale Price Indices with Regression Line")
    plt.xlabel("Sale Price Index")
    plt.ylabel("Total Quantity")

    # 显示图例
    plt.legend()

    # 显示网格
    plt.grid(True)

    # 显示图像
    plt.show()

    # 获取回归系数（a）和截距（b）
    a = model.coef_[0][0]
    b = model.intercept_[0]

    return a


if __name__ == "__main__":
    pass
    # dtype = [("type", "U10"), ("E", "f8"), ("R^2", "f8"), ("p_value", "f8")]
    # result = np.empty(6, dtype=dtype)
    # for i in range(1, 7):
    #     data = tool.get_np(f"problem_2/{i}.csv")
    #     a, b, c = log_log_demand_model_show(data)
    #     result[i - 1] = (mapping.get_typename(i), a, b, c)
    # tool.get_csv(result, "loglog_data.csv")
    i = 6
    data = tool.get_np(f"problem3/{101}.csv")
    print(mapping.get_name(101))
    filterd = filter_data2(data)
    linear_demand_get_param(filterd)
